from _dynamic_dataset.EvalType import EvalType
from _dynamic_dataset.ResultRecorder import ResultRecorder
from torch.utils.data import Dataset
from torchvision.transforms import ToTensor
import numpy as np
from PIL import Image
import torch


class EvalDataset(Dataset):
    """
    data：图像的数据集
    labels：标签
    maxNum：滑动窗口的最大大小
    transform：转换器
    """

    def __init__(self, data, labels, maxNum, threshold, transform=None):
        self.data = data
        self.labels = labels
        self.threshold = threshold
        self.transform = transform if transform is not None else ToTensor()
        # 构造判别结果记录器
        self.maxNum = maxNum
        self.resultRecord = ResultRecorder(maxNum, len(labels), self.threshold)

    def __getitem__(self, index):
        img = self.data[index]
        label = self.labels[index]
        # 如果图像是numpy.ndarray，先转换为PIL.Image
        if isinstance(img, np.ndarray):
            img = Image.fromarray(img)
        # 应用transform
        img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.data)



    # 对于数据集的添加子数据集的判定结果,默认全部为正确
    def addNewEvalRecord(self, indices, judge_result):
        # 默认不在子数据集判定的结果全为正确
        corrRecord = torch.ones((len(self.labels)), dtype=torch.int8)
        for i in range(len(indices)):
            corrRecord[indices[i]] = judge_result[i]
        self.resultRecord.addNewCorrRecord(corrRecord)

    # 对于没有进行过测试的结果，置信度延续上一轮的
    def addConfidencesRecord(self, indices, confidence_result, beta):
        """

        :param indices:
        :param confidence_result:
        :param beta: 置信度衰减系数
        :return:
        """
        last_confidences = self.resultRecord.getLatestConfidenceRecord()
        if beta != 1.0:
            last_confidences = last_confidences * beta
        for i in range(len(indices)):
            last_confidences[indices[i]] = confidence_result[i]
        self.resultRecord.addNewConfidenceRecord(last_confidences)

    def getIndicesFromEvalDataset(self, get_type, sample_ratio=None):
        if get_type == EvalType.NORMAL:
            return self.resultRecord.getHardIndicesFromCorrTensor()
        elif get_type == EvalType.RANDOM_CHOOSE_IGNORE_SET:
            return self.resultRecord.getIndicesWithRandomChoose(sample_ratio)
        elif get_type == EvalType.WEIGHT_CHOOSE_IGNORE_SET:
            return self.resultRecord.getIndicesWithWeightRandomChoose(sample_ratio)
        # elif get_type == EvalType.ROUND_NORMAL:
        #     return self.resultRecord.getIndicesFromCorrTensorWindows()